# uamcf/utils/visualizer.py
import os
from typing import Dict, Any, List, Tuple, Optional, Union

from .logger import get_logger

logger = get_logger("Visualizer")


def plot_compression_results(results: Dict[str, Any], output_dir: str = ".",
                             prefix: str = "compression") -> bool:
    """
    可视化压缩结果

    Args:
        results: 压缩结果字典
        output_dir: 输出目录
        prefix: 文件名前缀

    Returns:
        是否成功生成可视化
    """
    try:
        import matplotlib.pyplot as plt
        import numpy as np

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)

        # 提取原始和压缩后的指标
        if 'original' not in results or 'compressed' not in results:
            logger.error("Results must contain 'original' and 'compressed' data")
            return False

        original = results['original']
        compressed = results['compressed']
        improvements = results.get('improvements', {})

        # 1. 模型大小对比
        if 'size_bytes' in original and 'size_bytes' in compressed:
            plt.figure(figsize=(10, 6))

            # 转换为MB便于显示
            orig_size = original['size_bytes'] / (1024 * 1024)
            comp_size = compressed['size_bytes'] / (1024 * 1024)

            labels = ['Original', 'Compressed']
            sizes = [orig_size, comp_size]

            plt.bar(labels, sizes, color=['#3498db', '#2ecc71'])
            plt.title('Model Size Comparison')
            plt.ylabel('Size (MB)')
            plt.grid(axis='y', linestyle='--', alpha=0.7)

            # 添加数值标签
            for i, v in enumerate(sizes):
                plt.text(i, v + 0.1, f"{v:.2f} MB", ha='center')

            # 添加改进率
            if 'size_reduction' in improvements:
                reduction = improvements['size_reduction']
                plt.text(0.5, max(sizes) * 0.5, f"{reduction:.2f}x reduction",
                         ha='center', fontsize=14, bbox=dict(facecolor='white', alpha=0.5))

            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_size_comparison.png"), dpi=300)
            plt.close()

        # 2. 参数数量对比
        if 'param_count' in original and 'param_count' in compressed:
            plt.figure(figsize=(10, 6))

            # 转换为M便于显示
            orig_params = original['param_count'] / 1000000
            comp_params = compressed['param_count'] / 1000000

            labels = ['Original', 'Compressed']
            params = [orig_params, comp_params]

            plt.bar(labels, params, color=['#3498db', '#2ecc71'])
            plt.title('Model Parameter Count Comparison')
            plt.ylabel('Parameters (M)')
            plt.grid(axis='y', linestyle='--', alpha=0.7)

            # 添加数值标签
            for i, v in enumerate(params):
                plt.text(i, v + 0.1, f"{v:.2f}M", ha='center')

            # 添加改进率
            if 'param_reduction' in improvements:
                reduction = improvements['param_reduction']
                plt.text(0.5, max(params) * 0.5, f"{reduction:.2f}x reduction",
                         ha='center', fontsize=14, bbox=dict(facecolor='white', alpha=0.5))

            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_param_comparison.png"), dpi=300)
            plt.close()

        # 3. 精度对比
        if 'accuracy' in original and 'accuracy' in compressed:
            plt.figure(figsize=(10, 6))

            orig_acc = original['accuracy'] * 100  # 转为百分比
            comp_acc = compressed['accuracy'] * 100

            labels = ['Original', 'Compressed']
            accuracies = [orig_acc, comp_acc]

            bars = plt.bar(labels, accuracies, color=['#3498db', '#e74c3c' if comp_acc < orig_acc else '#2ecc71'])
            plt.title('Model Accuracy Comparison')
            plt.ylabel('Accuracy (%)')
            plt.ylim([min(min(accuracies) - 5, 90), 100])  # 设置合适的y轴范围
            plt.grid(axis='y', linestyle='--', alpha=0.7)

            # 添加数值标签
            for i, v in enumerate(accuracies):
                plt.text(i, v + 0.3, f"{v:.2f}%", ha='center')

            # 添加变化
            if 'accuracy_change' in improvements:
                change = improvements['accuracy_change'] * 100
                color = 'red' if change < 0 else 'green'
                plt.text(0.5, min(accuracies) - 2, f"Change: {change:+.2f}%",
                         ha='center', fontsize=14, color=color,
                         bbox=dict(facecolor='white', alpha=0.5))

            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_accuracy_comparison.png"), dpi=300)
            plt.close()

        # 4. 延迟对比
        if 'latency_ms' in original and 'latency_ms' in compressed:
            plt.figure(figsize=(10, 6))

            orig_latency = original['latency_ms']
            comp_latency = compressed['latency_ms']

            labels = ['Original', 'Compressed']
            latencies = [orig_latency, comp_latency]

            plt.bar(labels, latencies, color=['#3498db', '#2ecc71' if comp_latency < orig_latency else '#e74c3c'])
            plt.title('Model Inference Latency Comparison')
            plt.ylabel('Latency (ms)')
            plt.grid(axis='y', linestyle='--', alpha=0.7)

            # 添加数值标签
            for i, v in enumerate(latencies):
                plt.text(i, v + 0.1, f"{v:.2f} ms", ha='center')

            # 添加改进率
            if 'latency_improvement' in improvements:
                improvement = improvements['latency_improvement']
                plt.text(0.5, max(latencies) * 0.5, f"{improvement:.2f}x faster",
                         ha='center', fontsize=14, bbox=dict(facecolor='white', alpha=0.5))

            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_latency_comparison.png"), dpi=300)
            plt.close()

        # 5. 综合指标雷达图
        metrics = ['Size', 'Parameters', 'Speed', 'Accuracy']
        values = []

        if 'size_reduction' in improvements:
            values.append(min(improvements['size_reduction'], 5))  # 限制在5倍以内
        else:
            metrics.remove('Size')

        if 'param_reduction' in improvements:
            values.append(min(improvements['param_reduction'], 5))
        else:
            metrics.remove('Parameters')

        if 'latency_improvement' in improvements:
            values.append(min(improvements['latency_improvement'], 5))
        else:
            metrics.remove('Speed')

        if 'accuracy_ratio' in improvements:
            values.append(min(improvements['accuracy_ratio'], 1.2))  # 精度通常在1左右
        else:
            metrics.remove('Accuracy')

        if metrics and values:
            plt.figure(figsize=(8, 8))

            # 雷达图需要闭合的数据
            values_closed = np.append(values, values[0])
            metrics_closed = np.append(metrics, metrics[0])

            # 计算角度
            angles = np.linspace(0, 2 * np.pi, len(metrics), endpoint=False)
            angles_closed = np.append(angles, angles[0])

            # 绘制雷达图
            plt.polar(angles_closed, values_closed, marker='o')
            plt.fill(angles_closed, values_closed, alpha=0.25)

            # 添加标签
            plt.xticks(angles, metrics)

            # 添加参考线
            plt.yticks([1, 2, 3, 4, 5], ['1x', '2x', '3x', '4x', '5x'], color='grey')
            plt.ylim(0, 5)

            plt.title('Compression Improvements')
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_improvements_radar.png"), dpi=300)
            plt.close()

        # 6. 如果有步骤信息，绘制逐步压缩效果图
        if 'steps' in results and results['steps']:
            steps = results['steps']
            step_nums = list(range(1, len(steps) + 1))

            # 6.1 逐步精度变化
            accuracies = []
            for step in steps:
                if 'stats' in step and 'accuracy' in step['stats']:
                    accuracies.append(step['stats']['accuracy'] * 100)  # 转为百分比

            if accuracies:
                plt.figure(figsize=(12, 6))
                plt.plot(step_nums, accuracies, marker='o', linestyle='-', color='#3498db')
                plt.axhline(y=original['accuracy'] * 100, color='r', linestyle='--',
                            label=f"Original: {original['accuracy'] * 100:.2f}%")
                plt.title('Accuracy Change During Compression')
                plt.xlabel('Compression Step')
                plt.ylabel('Accuracy (%)')
                plt.legend()
                plt.grid(True, linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_step_accuracy.png"), dpi=300)
                plt.close()

            # 6.2 逐步大小变化
            sizes = []
            for step in steps:
                if 'stats' in step and 'size_bytes' in step['stats']:
                    sizes.append(step['stats']['size_bytes'] / (1024 * 1024))  # 转为MB

            if sizes:
                plt.figure(figsize=(12, 6))
                plt.plot(step_nums, sizes, marker='o', linestyle='-', color='#2ecc71')
                plt.axhline(y=original['size_bytes'] / (1024 * 1024), color='r', linestyle='--',
                            label=f"Original: {original['size_bytes'] / (1024 * 1024):.2f} MB")
                plt.title('Model Size Change During Compression')
                plt.xlabel('Compression Step')
                plt.ylabel('Size (MB)')
                plt.legend()
                plt.grid(True, linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_step_size.png"), dpi=300)
                plt.close()

        logger.info(f"Compression results visualization saved to {output_dir}")
        return True

    except ImportError:
        logger.error("Matplotlib is required for visualization. Install with 'pip install matplotlib'")
        return False
    except Exception as e:
        logger.error(f"Failed to create visualization: {str(e)}")
        return False


def plot_layer_statistics(model_analysis: Dict[str, Any], output_dir: str = ".",
                          prefix: str = "model_analysis") -> bool:
    """
    可视化模型层统计信息

    Args:
        model_analysis: 模型分析结果字典
        output_dir: 输出目录
        prefix: 文件名前缀

    Returns:
        是否成功生成可视化
    """
    try:
        import matplotlib.pyplot as plt
        import numpy as np

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)

        # 提取层信息
        if 'layers' not in model_analysis:
            logger.error("Model analysis must contain 'layers' data")
            return False

        layers = model_analysis['layers']

        # 1. 参数分布柱状图
        if layers:
            # 收集每一层的参数数量
            layer_names = []
            param_counts = []

            # 只选择参数量最大的前15层，避免图过于拥挤
            sorted_layers = sorted(layers, key=lambda x: x.get('parameters', 0), reverse=True)[:15]

            for layer in sorted_layers:
                name = layer.get('name', 'Unknown')
                # 截断过长的名称
                if len(name) > 20:
                    name = name[:17] + '...'
                layer_names.append(name)
                param_counts.append(layer.get('parameters', 0) / 1000)  # 转为K

            plt.figure(figsize=(12, 6))
            bars = plt.bar(layer_names, param_counts, color='#3498db')
            plt.title('Layer Parameter Distribution (Top 15 Layers)')
            plt.xlabel('Layer Name')
            plt.ylabel('Parameters (K)')
            plt.xticks(rotation=45, ha='right')
            plt.grid(axis='y', linestyle='--', alpha=0.7)
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_layer_parameters.png"), dpi=300)
            plt.close()

        # 2. 计算时间分布
        if 'computation_times' in model_analysis:
            comp_times = model_analysis['computation_times']

            if comp_times:
                # 只选择计算时间最长的前15层
                sorted_times = sorted(comp_times.items(), key=lambda x: x[1], reverse=True)[:15]
                layer_names = []
                times = []

                for name, time_ms in sorted_times:
                    # 截断过长的名称
                    if len(name) > 20:
                        name = name[:17] + '...'
                    layer_names.append(name)
                    times.append(time_ms)

                plt.figure(figsize=(12, 6))
                bars = plt.bar(layer_names, times, color='#e74c3c')
                plt.title('Layer Computation Time (Top 15 Layers)')
                plt.xlabel('Layer Name')
                plt.ylabel('Time (ms)')
                plt.xticks(rotation=45, ha='right')
                plt.grid(axis='y', linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_layer_computation.png"), dpi=300)
                plt.close()

        # 3. 内存使用分布
        if 'memory_stats' in model_analysis:
            memory_stats = model_analysis['memory_stats']

            if memory_stats:
                # 只选择内存使用最多的前15层
                sorted_memory = sorted(memory_stats.items(), key=lambda x: x[1]['total_size'], reverse=True)[:15]
                layer_names = []
                memory_sizes = []

                for name, stats in sorted_memory:
                    # 截断过长的名称
                    if len(name) > 20:
                        name = name[:17] + '...'
                    layer_names.append(name)
                    memory_sizes.append(stats['total_size'] / (1024 * 1024))  # 转为MB

                plt.figure(figsize=(12, 6))
                bars = plt.bar(layer_names, memory_sizes, color='#f39c12')
                plt.title('Layer Memory Usage (Top 15 Layers)')
                plt.xlabel('Layer Name')
                plt.ylabel('Memory (MB)')
                plt.xticks(rotation=45, ha='right')
                plt.grid(axis='y', linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_layer_memory.png"), dpi=300)
                plt.close()

        # 4. 瓶颈图
        if 'bottlenecks' in model_analysis:
            bottlenecks = model_analysis['bottlenecks']

            # 绘制计算瓶颈
            if 'compute' in bottlenecks and bottlenecks['compute']:
                compute_bottlenecks = bottlenecks['compute'][:10]  # 最多10个
                layer_names = []
                ratios = []

                for b in compute_bottlenecks:
                    name = b.get('name', 'Unknown')
                    # 截断过长的名称
                    if len(name) > 20:
                        name = name[:17] + '...'
                    layer_names.append(name)
                    ratios.append(b.get('ratio', 0) * 100)  # 转为百分比

                plt.figure(figsize=(12, 6))
                bars = plt.bar(layer_names, ratios, color='#9b59b6')
                plt.title('Computation Bottlenecks')
                plt.xlabel('Layer Name')
                plt.ylabel('Percentage of Total Computation (%)')
                plt.xticks(rotation=45, ha='right')
                plt.grid(axis='y', linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_compute_bottlenecks.png"), dpi=300)
                plt.close()

            # 绘制内存瓶颈
            if 'memory' in bottlenecks and bottlenecks['memory']:
                memory_bottlenecks = bottlenecks['memory'][:10]  # 最多10个
                layer_names = []
                ratios = []

                for b in memory_bottlenecks:
                    name = b.get('name', 'Unknown')
                    # 截断过长的名称
                    if len(name) > 20:
                        name = name[:17] + '...'
                    layer_names.append(name)
                    ratios.append(b.get('ratio', 0) * 100)  # 转为百分比

                plt.figure(figsize=(12, 6))
                bars = plt.bar(layer_names, ratios, color='#1abc9c')
                plt.title('Memory Bottlenecks')
                plt.xlabel('Layer Name')
                plt.ylabel('Percentage of Total Memory (%)')
                plt.xticks(rotation=45, ha='right')
                plt.grid(axis='y', linestyle='--', alpha=0.7)
                plt.tight_layout()
                plt.savefig(os.path.join(output_dir, f"{prefix}_memory_bottlenecks.png"), dpi=300)
                plt.close()

        # 5. 层类型分布饼图
        if layers:
            layer_types = {}
            for layer in layers:
                layer_type = layer.get('type', 'Unknown')
                if layer_type in layer_types:
                    layer_types[layer_type] += 1
                else:
                    layer_types[layer_type] = 1

            plt.figure(figsize=(10, 8))
            plt.pie(layer_types.values(), labels=layer_types.keys(), autopct='%1.1f%%',
                    startangle=90, shadow=True)
            plt.title('Layer Type Distribution')
            plt.axis('equal')  # 保持圆形
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, f"{prefix}_layer_types.png"), dpi=300)
            plt.close()

        logger.info(f"Model analysis visualization saved to {output_dir}")
        return True

    except ImportError:
        logger.error("Matplotlib is required for visualization. Install with 'pip install matplotlib'")
        return False
    except Exception as e:
        logger.error(f"Failed to create visualization: {str(e)}")
        return False